import re
import os
import sys
import datetime


def get_arg(prefix):
    py_path = os.path.join(prefix, 'megatron/training/arguments.py')
    with open(py_path, 'r') as f:
        lines = f.readlines()
    lines = [x for x in lines if '.add_argument(' in x]
    args = []
    pattern = r""".+.add_argument\(["'](--[^,]+)['"],[^,]*"""
    for line in lines:
        matches = re.search(pattern, line)
        if not matches:
            print(f'{line} not find arguments.')
        else:
            args.append(matches.group(1))
    return args


def get_args_chgs(prefix_last, prefix, save_path):
    args_last = get_arg(prefix_last)
    args = get_arg(prefix)
    adds = [x for x in args if x not in args_last]
    dels = [x for x in args_last if x not in args]
    adds_str = '\n' + '\n'.join(adds).strip() if len(adds) else ''
    dels_str = '\n' + '\n'.join(dels).strip() if len(dels) > 0 else ''
    save_str = (f"{datetime.date.today().strftime('%Y%m%d')}\n"
                f"{'{:-<50}'.format(str(len(adds))+ ' adds:')}{adds_str}\n"
                f"{'{:-<51}'.format(str(len(dels))+ ' dels:')}{dels_str}")
    savedir = os.path.dirname(save_path)
    if not savedir == '' and not os.path.exists(savedir):
        os.makedirs(savedir)
    with open(save_path, 'w') as f:
        f.write(save_str)


def main():
    mg_path_last = sys.argv[1]
    mg_path = sys.argv[2]
    save_path = sys.argv[3]
    get_args_chgs(mg_path_last, mg_path, save_path)


if __name__ == "__main__":
    main()
